package com.wzw.springboot.principle.use.websock.second;

import com.wzw.springboot.principle.use.security.infrastructure.model.SecurityUser;
import com.wzw.springboot.principle.use.security.infrastructure.provider.QrCodeAuthenticationToken;
import com.wzw.springboot.principle.use.security.infrastructure.provider.SmsCodeAuthenticationToken;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.security.Principal;
import java.util.Map;
import java.util.Objects;

/**
 * WebSocket 握手拦截器
 *
 * @author Wangzhiwen
 */
@Component
@SuppressWarnings("NullableProblems")
public class PrincipleHandshakeInterceptor implements HandshakeInterceptor {

    /**
     * 握手前执行拦截 认证
     *
     * @param request    请求体
     * @param response   响应体
     * @param wsHandler  拦截执行体
     * @param attributes 属性键值对
     * @return true 放行 false 不放行
     * @throws Exception 异常
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // spring security
        Principal principal = request.getPrincipal();
        if (Objects.isNull(principal)) {
            return false;
        }
        String key = SecurityUser.class.getSimpleName();
        if (principal instanceof UsernamePasswordAuthenticationToken) {
            UsernamePasswordAuthenticationToken usernamePasswordAuthenticationToken = (UsernamePasswordAuthenticationToken) principal;
            attributes.put(key, usernamePasswordAuthenticationToken.getPrincipal());
        } else if (principal instanceof QrCodeAuthenticationToken) {
            QrCodeAuthenticationToken qrCodeAuthenticationToken = (QrCodeAuthenticationToken) principal;
            attributes.put(key, qrCodeAuthenticationToken.getPrincipal());
        } else if (principal instanceof SmsCodeAuthenticationToken) {
            SmsCodeAuthenticationToken smsCodeAuthenticationToken = (SmsCodeAuthenticationToken) principal;
            attributes.put(key, smsCodeAuthenticationToken.getPrincipal());
        } else {
            return false;
        }
        return true;
    }

    /**
     * 握手后拦截
     *
     * @param request   请求体
     * @param response  响应体
     * @param wsHandler 拦截执行体
     * @param exception 异常
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
    }
}
